# 给你一个 n 个点的带权无向连通图，节点编号为 0 到 n-1 ，同时还有一个数组 edges ，
# 其中 edges[i] = [fromi, toi, weighti] 表示在 fromi 和 toi 节点之间有一条带权无向边。
# 最小生成树 (MST) 是给定图中边的一个子集，它连接了所有节点且没有环，而且这些边的权值和最小。
#
#  请你找到给定图中最小生成树的所有关键边和伪关键边。
#  如果从图中删去某条边，会导致最小生成树的权值和增加，那么我们就说它是一条关键边。
#  伪关键边则是可能会出现在某些最小生成树中但不会出现在所有最小生成树中的边。
#
#  请注意，你可以分别以任意顺序返回关键边的下标和伪关键边的下标。
#
#  示例 1：
#  输入：n = 5, edges = [[0,1,1],[1,2,1],[2,3,2],[0,3,2],[0,4,3],[3,4,3],[1,4,6]]
# 输出：[[0,1],[2,3,4,5]]
# 解释：上图描述了给定图。
# 下图是所有的最小生成树。
#
# 注意到第 0 条边和第 1 条边出现在了所有最小生成树中，所以它们是关键边，我们将这两个下标作为输出的第一个列表。
# 边 2，3，4 和 5 是所有 MST 的剩余边，所以它们是伪关键边。我们将它们作为输出的第二个列表。
#
#  示例 2 ：
#  输入：n = 4, edges = [[0,1,1],[1,2,1],[2,3,1],[0,3,1]]
# 输出：[[],[0,1,2,3]]
# 解释：可以观察到 4 条边都有相同的权值，任选它们中的 3 条可以形成一棵 MST 。所以 4 条边都是伪关键边。
import collections
import heapq
import sys
from typing import List


class UnionFindSet:
    def __init__(self, n: int):
        self.parent = [i for i in range(n)]
        # self.rank = [1] * n
        self.circlesNum = n

    def find(self, node: int) -> int:
        if self.parent[node] != node:
            self.parent[node] = self.find(self.parent[node])
        return self.parent[node]

    def union(self, node1: int, node2: int) -> bool:
        p1, p2 = self.find(node1), self.find(node2)
        if p1 == p2:
            return False
        else:
            self.circlesNum -= 1
            self.parent[p2] = p1
            return True


class Solution:
    def findCriticalAndPseudoCriticalEdges(self, n: int, edges: List[List[int]]) -> List[List[int]]:
        """
        先通过kruskal 或者 prim算法算出MST的总权值和
        然后
        遍历所有的边执行如下操作：
            在删掉某一条后边再重新运行kruskal算法求最小生成树的权值和(prim算法也可以实现)
            如果权值和变大则是关键边

            在加入某一条边后再重新运行kruskal算法求最小生成树的权值和(prim算法也可以实现)
            如果不变则说明是伪关键边
        :param n:
        :param edges:
        :return:
        """
        sortedEdges = sorted(edges, key=lambda e: e[2])
        keyEdges, pseuEdges = [], []

        def kruskalWithSomeDeal(index: int, giveOrTake: bool) -> int:
            """
            对kruskal算法进行改造(加入关建边和伪关键边的处理)
            :param index: 要预先加入最小生成树的边或者去掉边在初始数组的索引(-1表示既不加入边也不删掉边)
            :param giveOrTake: 提前给定一条边时为false，提前去掉一条边为true
            :return:
            """
            mstWeightSum = 0
            uf = UnionFindSet(n)
            # 伪关键边处理，先将边加入到MST中(即 将边的两个顶点合并到一个集合中)
            if not giveOrTake and index >= 0:
                fromm, to, weight = edges[index]
                uf.union(fromm, to)
                mstWeightSum += weight
            for ii, tmp in enumerate(sortedEdges):
                fromm, to, weight = tmp
                if index >= 0 and giveOrTake and sortedEdges[ii] == edges[index]:  # 关键边处理(去掉某条边)
                    continue
                if uf.union(fromm, to):
                    mstWeightSum += weight
            return mstWeightSum if uf.circlesNum == 1 else sys.maxsize

        mstWeightSum = kruskalWithSomeDeal(-1, False)  # 先通过kruskal/prim 算法算出MST的总权值和
        # print(mstWeightSum)
        for i in range(len(edges)):  # 对每条边都分别进行判断是属于关建边还是伪关键边
            if kruskalWithSomeDeal(i, True) > mstWeightSum:  # 删掉某一条后边再重新运行kruskal算法求最小生成树的权值和
                keyEdges.append(i)
            elif kruskalWithSomeDeal(i, False) == mstWeightSum:  # 提前就加入某一条边后再重新运行kruskal算法求最小生成树的权值和
                pseuEdges.append(i)
        return [keyEdges, pseuEdges]

    def kruskal(self, n: int, edges: List[List[int]]) -> int:
        """
        普通的kruskal算法求MST(贪心 + 并查集):
        对边进行操作，常用于稀疏图
        将边按大小进行排序，
        然后每次从边集中取出权重最小且两个顶点都不在同一个集合的边加入生成树中(并查集实现)
        :param n:
        :param edges:
        :return:
        """
        sortedEdges = sorted(edges, key=lambda e: e[2])
        mstWeightSum = 0
        uf = UnionFindSet(n)
        for formm, to, weight in sortedEdges:
            if uf.union(formm, to):
                mstWeightSum += weight
        return mstWeightSum if uf.circlesNum == 1 else sys.maxsize

    def prim(self, n: int, edges: List[List[int]]) -> int:
        """
        对节点进行操作，常用于稠密图
        先随便找一个点，然后每次从已经遍历的点的所有边中找到一条边要满足如下条件，将 to的顶点加入到已遍历集合中
            1.该边的权值最小
            2.该边有个顶点未被遍历
        添加一个点
        :param n:
        :param sortedEdges:
        :return:
        """
        def getAdjList() -> dict:
            adjList = collections.defaultdict(list)
            for fromm, to, weight in edges:
                adjList[fromm].append((weight, to))
                adjList[to].append((weight, fromm))
            return adjList

        adjList = getAdjList()
        visited = set()
        queue = [0]
        minHeap = []  # 存放边的最小堆
        mstWeightSum = 0
        while queue:
            node = queue[0]
            del queue[0]
            if node not in visited and (neighbors := adjList[node]):
                visited.add(node)
                minHeap.extend(neighbors)
                heapq.heapify(minHeap)
                while minHeap:
                    weight, to = heapq.heappop(minHeap)
                    if to not in visited:
                        mstWeightSum += weight
                        queue.append(to)
                        break
        return mstWeightSum if len(visited) == n else sys.maxsize


if __name__ == "__main__":
    # n, edges = 5, [[0, 1, 1], [1, 2, 1], [2, 3, 2], [0, 3, 2], [0, 4, 3], [3, 4, 3], [1, 4, 6]]
    n, edges = 4, [[0, 1, 1], [1, 2, 1], [2, 3, 1], [0, 3, 1]]
    # print(Solution().prim(n, edges))
    # print(Solution().kruskal(n, edges))
    print(Solution().findCriticalAndPseudoCriticalEdges(n, edges))
